import json

__author__ = 'sdjkx'

import os
from libs.ustr import *
 
try:
  import xml.etree.cElementTree as ET
except ImportError:
  import xml.etree.ElementTree as ET
from PyQt5.QtWidgets import QMessageBox,QFileDialog

class ScanXmlSortLabel:

    def __init__(self, labelImg):
        self.labelImg = labelImg
        self.ImageWidth = 0

    def replaceSpace(self):
        sourcePath = "."

        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择XML文件路径", sourcePath,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirpath is not None and len(dirpath) > 1:
            ret = QMessageBox.question(self.labelImg, '提示', "确定要对目录【"+dirpath+"】里的XML文件的中的半角空格改为全角吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)
            if ret == QMessageBox.No:
                return

        newPath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择XML文件路径", dirpath,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))
        if newPath is None:
            QMessageBox.information(self.labelImg, '提示', "必须选择保存XML文件的路径！")
            return

        self.replaceAllSpace(dirpath, newPath)

    def replaceAllSpace(self, xmlPath, targetPath):
        xmlFileList = os.listdir(xmlPath)
        for xmlFileName in xmlFileList:
            xmlFilePath =  os.path.join(xmlPath, xmlFileName)
            newXmlFilePath = os.path.join(targetPath, xmlFileName)

            xmlText = ""
            for line in open(xmlFilePath,"r", encoding='utf-8'): #设置文件对象并读取每一行文件
                if line.find("<name>") > 0:
                    oldTxt = line.strip()
                    txt = "    " + oldTxt.replace(" ", "　")
                    xmlText += txt
                else:
                    xmlText += line

            self.saveXML(xmlText, newXmlFilePath)

    def scanCPLabelFun(self):
        sourcePath = "."

        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择XML文件路径", sourcePath,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirpath is not None and len(dirpath) > 1:
            ret = QMessageBox.question(self.labelImg, '提示', "确定要对目录【"+dirpath+"】里的XML文件的标签进行排序吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)
            if ret == QMessageBox.No:
                return

        newPath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择导出文件路径", dirpath,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))
        if newPath is None:
            QMessageBox.information(self.labelImg, '提示', "必须选择保存XML文件的路径！")
            return

        self.scanCpAndBusLabel(dirpath, newPath)

    def getXmlDataToJson(self,id,lableId, xmlFilePath):
        imageList = []

        tree = ET.parse(xmlFilePath)     # 打开xml文档
        root = tree.getroot()         #获得root节点
        fileName = root.find("filename").text
        # filePath = root.find("path").text
        size = root.find("size")
        w = int(size.find("width").text)
        h = int(size.find("height").text)

        try:
            for object in root.findall('object'): #找到root节点下的所有object节点
                lableId += 1
                bndbox = object.find('bndbox')      # 子节点下属性bndbox的值
                xmin = int(bndbox.find('xmin').text)
                ymin = int(bndbox.find('ymin').text)
                xmax = int(bndbox.find('xmax').text)
                ymax = int(bndbox.find('ymax').text)
                if object.find("category") is not None:
                    category = object.find("category").text
                else:
                    category = "文本"
                i = 0
                category_id = 1
                for txt in self.labelImg.label_hist:
                    i += 1
                    if txt == category:
                        category_id = i
                        break

                if xmin < 0:
                    xmin = 0
                if ymin < 0:
                    ymin = 0

                words = object.find("name").text
                dataObj = {
                    "segmentation":[],    # coco中分割相关的信息，默认填空即可
                    "area":100.0,         # 目前无需更改
                    "iscrowd":0,          # 目前无需更改
                    "image_id":id,         # images 中对应图片的id
                    "bbox":[              # 标注框 左上角点坐标(X, Y) 和 宽高，即 [X, Y , W, H]
                        xmin,
                        ymin,
                        xmax - xmin,
                        ymax - ymin
                    ],
                    "category_id":category_id,      # 标注框对应的 识别字段id，该id在 categories 中定义
                    "id":lableId,           # 标注框id，依次递增
                    "extra_infos":{
                        "labelValue":words   # 标注框中的内容
                    }
                }

                imageList.append(dataObj)
        except:
            print("解析文件失败：", xmlFilePath)
            fileName = ""

        return fileName,w, h , imageList,lableId

    def getXmlData(self, xmlFilePath):
        imageDict = {}
        imageList = [[] for i in range(3)]  # 创建的是多行三列的二维列表

        scale = 1
        tree = ET.parse(xmlFilePath)     # 打开xml文档
        root = tree.getroot()         #获得root节点
        fileName = root.find("filename").text
        filePath = root.find("path").text
        size = root.find("size")
        w = size.find("width").text
        h = size.find("height").text


        try:
            for object in root.findall('object'): #找到root节点下的所有object节点
                bndbox = object.find('bndbox')      # 子节点下属性bndbox的值
                xmin = float(bndbox.find('xmin').text)
                ymin = float(bndbox.find('ymin').text)
                xmax = float(bndbox.find('xmax').text)
                ymax = float(bndbox.find('ymax').text)

                left = int(xmin * scale)
                upper = int(ymin * scale)
                right = int(xmax * scale)
                lower = int(ymax * scale)
                if left < 0:
                    left = 0
                if upper < 0:
                    upper = 0

                key = str(left + right) + "-" + str(upper)
                # rowid = object.find("rowid").text

                words = object.find("name").text
                imageDict[key] = '  <object>\n' \
                                 '    <name>'+words+'</name>\n' \
                                 '    <pose>Unspecified</pose>\n' \
                                 '    <truncated>0</truncated>\n' \
                                 '    <difficult>0</difficult>\n' \
                                 '    <rowid>xxxxxxxxxx</rowid>\n' \
                                 '    <bndbox>\n' \
                                 '      <xmin>'+str(xmin)+'</xmin>\n' \
                                 '      <ymin>'+str(ymin)+'</ymin>\n' \
                                 '      <xmax>'+str(xmax)+'</xmax>\n' \
                                 '      <ymax>'+str(ymax)+'</ymax>\n' \
                                 '    </bndbox>\n' \
                                 '  </object>\n'
                if ymin < 300:
                    imageList[0].append([left + right, upper])
                elif ymin < 400:
                    imageList[1].append([left + right, upper])
                else:
                    imageList[2].append([left + right, upper])
        except:
            print("解析文件失败：", xmlFilePath)
            fileName = ""

        return fileName,filePath,w, h, imageDict, imageList

    def scanXMLToJsonFun(self):
        sourcePath = "."

        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择导出为JSON的XML文件路径", sourcePath,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirpath is not None and len(dirpath) > 1:
            ret = QMessageBox.question(self.labelImg, '提示', "确定要对目录【"+dirpath+"】里的XML文件的标签导出为OCR的JSON文件吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)
            if ret == QMessageBox.No:
                return

        self.scanXMLToJson(dirpath)


    def scanXMLToJson(self, xmlPath):
        train = {}
        images = []
        annotations = []
        categories = []

        train["images"] = images
        train["annotations"] = annotations
        train["categories"] = categories

        xmlFileList = os.listdir(xmlPath)
        id = 0
        lableId = 0

        # --------------categories--------------------------------
        categoryId = 0
        for item in self.labelImg.label_hist:
            categoryId += 1
            categorieObj = {
                "name":item,    # 识别字段名称
                "id":categoryId              # category_id
            }
            categories.append(categorieObj)

        print(categories)

        for xmlFileName in xmlFileList:
            id += 1
            xmlFilePath =  os.path.join(xmlPath, xmlFileName)

            print("解析的文件：", xmlFilePath)

            fileName,w, h, annotationList,lableId = self.getXmlDataToJson(id,lableId, xmlFilePath)

            if len(annotationList) > 0:
                # ---------------------images ---------------------------
                imageObj = {
                    "file_name": fileName,
                    "height":h,
                    "width":w,
                    "depth":3,
                    "id":id
                }
                images.append(imageObj)

                # ----------------------annotations-----------------------
                for annotation in annotationList:
                    annotations.append(annotation)

        # jsonStr = json.dumps(train)# , sort_keys=True, indent=4, separators=(',', ': '))
        print("XML解析完成，开始保存JSON到文件：",xmlPath)
        self.saveStr(train, xmlPath)


    def saveStr(self, train, filePath):
        path = os.path.join(filePath, "train.json")
        # content=str(jsonStr.encode('utf-8'),encoding='utf-8')
        with open(path,"w",encoding='utf-8') as f:
            json.dump(train,f,ensure_ascii=False,sort_keys=False, indent=4)
            print(u'JSON数据保存到文件完成...')

    # 获取列表的第二个元素
    def takeSecond(self,elem):
        return elem[0]

    def scanCpAndBusLabel(self, xmlPath, targetPath):
        xmlFileList = os.listdir(xmlPath)
        for xmlFileName in xmlFileList:
            xmlFilePath =  os.path.join(xmlPath, xmlFileName)
            newXmlFilePath = os.path.join(targetPath, xmlFileName)
            if os.path.exists(newXmlFilePath):
                continue

            print("解析的文件：", xmlFilePath)

            fileName,filePath,w, h, imageDict, imageList = self.getXmlData(xmlFilePath)
            if fileName == "":
                continue

            newImageList = []  # 创建的是多行三列的二维列表
            for list in imageList:
                list.sort(key=self.takeSecond, reverse = True)
                newImageList.append(list)
            # print(newImageList)

            rowid = 1
            objectList = ""
            for list in newImageList:
                for item in list:
                    key = str(item[0]) + '-' + str(item[1])
                    value = imageDict[key]
                    # print(value)
                    value = value.replace('xxxxxxxxxx', str(rowid))
                    objectList += value
                    rowid += 1

            xmlComment = "<annotation>\n" \
                 " <folder>train</folder>\n" \
                 " <filename>"+fileName+"</filename>\n" \
                 " <path>"+filePath+"</path>\n" \
                 " <source>\n" \
                 "     <database>Unknown</database>\n" \
                 " </source>\n" \
                 " <size>\n" \
                 "     <width>"+w+"</width>\n" \
                 "     <height>"+h+"</height>\n" \
                 "     <depth>3</depth>\n" \
                 " </size>\n" \
                 " <segmented>0</segmented>\n"

            xmlComment = xmlComment + objectList + "</annotation>"
            self.saveXML(xmlComment, newXmlFilePath)

    def saveXML(self, xml, filePath):
        content=str(xml.encode('utf-8'),encoding='utf-8')

        with open(filePath, 'a', encoding='utf-8') as file_handle:   # .txt可以不自己新建,代码会自动新建
            file_handle.write(content)     # 写入
            file_handle.write('\n')         # 有时放在循环里面需要自动转行，不然会覆盖上一条数据


